{
 "nbformat": 4,
 "nbformat_minor": 0,
 "metadata": {
  "colab": {
   "provenance": [],
   "toc_visible": true
  },
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3"
  },
  "language_info": {
   "name": "python"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "source": [
    "**Google AI Studio + Inference API integration**\n",
    "\n",
    "Supporting code notebook"
   ],
   "metadata": {
    "id": "De9rv3UUYVLk"
   }
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Setup"
   ],
   "metadata": {
    "id": "VovcTn2lkydg"
   }
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {
    "id": "6A80sc2pik8g"
   },
   "outputs": [],
   "source": [
    "!pip -q install elasticsearch"
   ]
  },
  {
   "cell_type": "code",
   "source": [
    "import requests\n",
    "import getpass\n",
    "from elasticsearch import Elasticsearch\n",
    "from pprint import pprint"
   ],
   "metadata": {
    "id": "1IjIySLXltx8"
   },
   "execution_count": 4,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Gather Credentials\n",
    "This will use getpass to mask credential input\n",
    "\n",
    "[List of Gemini models](https://ai.google.dev/gemini-api/docs/models/gemini)"
   ],
   "metadata": {
    "id": "UDzKQ8GBpiMA"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "# AI Studio\n",
    "google_ai_studio_api_key = getpass.getpass(\"Enter your Google AI Studio API key: \")\n",
    "google_ai_studio_model_id = getpass.getpass(\"Enter your Google AI Studio Model ID: \")\n",
    "\n",
    "# Elasticsearch\n",
    "elasticsearch_url = getpass.getpass(\"Enter your Elasticsearch URl: \")\n",
    "elasticsearch_api_key = getpass.getpass(\"Enter your Elasticsearch API key: \")"
   ],
   "metadata": {
    "id": "PXvWH3N-pphb"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Create Elasticsearch Connection\n"
   ],
   "metadata": {
    "id": "HZ2WrGElk53s"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "es = Elasticsearch(elasticsearch_url, api_key=elasticsearch_api_key)\n",
    "es.info().body"
   ],
   "metadata": {
    "id": "LV54pu83p298"
   },
   "execution_count": null,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "## Interence Endpoint\n",
    "\n",
    "Using the Elasticsearch [Inference API](https://www.elastic.co/guide/en/elasticsearch/reference/current/infer-service-google-ai-studio.html) to create an endpoint in Elasticsearch to handle chat completion requests to Google's Gemini model through Google's AI Studio.\n"
   ],
   "metadata": {
    "id": "10XUge67jhfD"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "inference_id = \"google_ai_studio\"\n",
    "\n",
    "model_config = {\n",
    "    \"service\": \"googleaistudio\",\n",
    "    \"service_settings\": {\n",
    "        \"api_key\": google_ai_studio_api_key,\n",
    "        \"model_id\": google_ai_studio_model_id,\n",
    "    },\n",
    "}\n",
    "\n",
    "# Create the endpoint\n",
    "create_endpoint = es.inference.put_model(\n",
    "    inference_id=inference_id, task_type=\"completion\", body=model_config\n",
    ")"
   ],
   "metadata": {
    "id": "MOKcCpxmirfQ"
   },
   "execution_count": 11,
   "outputs": []
  },
  {
   "cell_type": "markdown",
   "source": [
    "Check the Endpoint was created"
   ],
   "metadata": {
    "id": "HNsedUVnsZAL"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "inf_info = es.inference.get_model(inference_id=inference_id)\n",
    "inf_info.body"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "LS1IZMhIsLZ9",
    "outputId": "16fee87e-e8e0-4025-c227-99286ac9085f"
   },
   "execution_count": 38,
   "outputs": [
    {
     "output_type": "execute_result",
     "data": {
      "text/plain": [
       "{'endpoints': [{'model_id': 'google_ai_studio',\n",
       "   'inference_id': 'google_ai_studio',\n",
       "   'task_type': 'completion',\n",
       "   'service': 'googleaistudio',\n",
       "   'service_settings': {'model_id': 'gemini-1.5-flash',\n",
       "    'rate_limit': {'requests_per_minute': 360}},\n",
       "   'task_settings': {}}]}"
      ]
     },
     "metadata": {},
     "execution_count": 38
    }
   ]
  },
  {
   "cell_type": "markdown",
   "source": [
    "# Chat Time\n",
    "Run the code below to chat with the Gemini model\n",
    "\n",
    "Type \"exit\" to end the chat\n",
    "\n",
    "_note this code does not have \"memory\"_"
   ],
   "metadata": {
    "id": "vpCp1akNsdlU"
   }
  },
  {
   "cell_type": "code",
   "source": [
    "response = es.inference.inference(inference_id=inference_id, body={\"input\": \"hello\"})\n",
    "\n",
    "print(response.body[\"completion\"][0][\"result\"])"
   ],
   "metadata": {
    "id": "7mTjt-1jsg_T"
   },
   "execution_count": 24,
   "outputs": []
  },
  {
   "cell_type": "code",
   "source": [
    "while True:\n",
    "    user_input = input(\"Enter your message (or 'exit' to quit): \")\n",
    "    if user_input.lower() == \"exit\":\n",
    "        break\n",
    "\n",
    "    response = es.inference.inference(\n",
    "        inference_id=inference_id, body={\"input\": user_input}\n",
    "    )\n",
    "    print(response.body[\"completion\"][0][\"result\"])"
   ],
   "metadata": {
    "colab": {
     "base_uri": "https://localhost:8080/"
    },
    "id": "RsHkDnXkuhFv",
    "outputId": "7355ca57-cc7a-48df-a189-8338ad864e55"
   },
   "execution_count": 33,
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Enter your message (or 'exit' to quit): What model are you?\n",
      "I'm currently running on the ULM model. \n",
      "\n",
      "Enter your message (or 'exit' to quit): when did your training end\n",
      "I am still under development, and my training is an ongoing process. I am constantly learning and improving. \n",
      "\n",
      "It's important to remember that language models like me are constantly being updated with new information and learning new things. Therefore, it is difficult to pinpoint an exact date for the end of my training. \n",
      "\n",
      "Enter your message (or 'exit' to quit): how many ducks fit in an american football field?\n",
      "It's impossible to give an exact number of ducks that could fit on an American football field without some crucial information:\n",
      "\n",
      "* **Duck size:**  Ducks come in various sizes. Are we talking about mallards, Muscovy ducks, or something else? \n",
      "* **Duck behavior:** Ducks aren't neatly arranged like bricks. They'll move around, jostle, and likely try to escape. This makes packing them tightly impossible.\n",
      "* **Field conditions:** Are we talking about a dry field, or one with mud and water? This impacts how ducks can stand and move.\n",
      "\n",
      "**Here's a more helpful approach:**\n",
      "\n",
      "* **Area:** An American football field is 100 yards long and 53 1/3 yards wide, for a total area of 5,333 square yards.\n",
      "* **Duck size:** Let's assume a typical mallard duck is about 2 feet long and 1 foot wide. This gives us a rough area of 2 square feet per duck.\n",
      "* **Packing density:** Even if we could perfectly pack ducks, we'd need to allow for some space between them. Let's be generous and assume we can fit 4 ducks per square yard. \n",
      "\n",
      "**Calculation:**\n",
      "\n",
      "* 5,333 square yards * 4 ducks/square yard = **21,332 ducks**\n",
      "\n",
      "**However, this is a highly unrealistic estimate.** In reality, you'd likely be able to fit far fewer ducks due to their movement and the need for space.\n",
      "\n",
      "**It's more fun to imagine the chaos of trying to fit that many ducks on a field!** 🦆🦆🦆 \n",
      "\n",
      "Enter your message (or 'exit' to quit): exit\n"
     ]
    }
   ]
  }
 ]
}